package org.example.cases.sf1;

import com.antgroup.geaflow.api.function.base.FilterFunction;
import com.antgroup.geaflow.api.function.base.MapFunction;
import com.antgroup.geaflow.api.graph.PGraphWindow;
import com.antgroup.geaflow.api.graph.compute.VertexCentricCompute;
import com.antgroup.geaflow.api.graph.function.vc.VertexCentricCombineFunction;
import com.antgroup.geaflow.api.graph.function.vc.VertexCentricComputeFunction;
import com.antgroup.geaflow.api.pdata.stream.window.PWindowSource;
import com.antgroup.geaflow.api.pdata.stream.window.PWindowStream;
import com.antgroup.geaflow.api.window.impl.AllWindow;
import com.antgroup.geaflow.common.tuple.Tuple;
import com.antgroup.geaflow.model.graph.edge.EdgeDirection;
import com.antgroup.geaflow.model.graph.edge.IEdge;
import com.antgroup.geaflow.model.graph.edge.impl.ValueEdge;
import com.antgroup.geaflow.model.graph.vertex.IVertex;
import com.antgroup.geaflow.model.graph.vertex.impl.ValueVertex;
import com.antgroup.geaflow.pipeline.Pipeline;
import com.antgroup.geaflow.pipeline.task.PipelineTask;
import com.antgroup.geaflow.view.GraphViewBuilder;
import com.antgroup.geaflow.view.IViewDesc;
import com.antgroup.geaflow.view.graph.GraphViewDesc;
import org.example.constants.Constants;
import org.example.function.AbstractVcFunc;
import org.example.function.CSVFileSinkV2;
import org.example.function.CSVFileSourceV2;
import org.example.model.LongStringResult;
import org.example.util.FileUtils;

import java.util.Arrays;
import java.util.Collections;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Set;

public class Case1Pipeline {

    private final String snapshotPath;

    private final Pipeline pipeline;

    public Case1Pipeline(String snapshotPath, Pipeline pipeline) {
        this.snapshotPath = snapshotPath;
        this.pipeline = pipeline;
    }

    public void submit() {

        pipeline.submit((PipelineTask) pipelineTaskCxt -> {
            PWindowSource<IVertex<Long, Double>> accountVertices = pipelineTaskCxt.buildSource(new CSVFileSourceV2<>(
                            FileUtils.pathJoin(this.snapshotPath, Constants.FILE_ACCOUNT),
                            Collections.singletonList(Constants.ACCOUNT_ID),
                            (CSVFileSourceV2.CSVLineParser<IVertex<Long, Double>>) line -> {
                                IVertex<Long, Double> vertex = new ValueVertex<>(
                                        Long.valueOf(line[0]),
                                        null
                                );
                                return Collections.singletonList(vertex);
                            }
                    ), AllWindow.getInstance())
                    .withParallelism(Constants.SOURCE_PARALLELISM);

            PWindowSource<IVertex<Long, Double>> loanVertices = pipelineTaskCxt.buildSource(new CSVFileSourceV2<>(
                            FileUtils.pathJoin(this.snapshotPath, Constants.FILE_LOAN),
                            Arrays.asList(Constants.LOAN_ID, Constants.LOAN_AMOUNT),
                            (CSVFileSourceV2.CSVLineParser<IVertex<Long, Double>>) line -> {
                                IVertex<Long, Double> vertex = new ValueVertex<>(
                                        -Long.parseLong(line[0]),
                                        Double.valueOf(line[1])
                                );
                                return Collections.singletonList(vertex);
                            }
                    ), AllWindow.getInstance())
                    .withParallelism(Constants.SOURCE_PARALLELISM);

            PWindowSource<IVertex<Long, Double>> personVertices = pipelineTaskCxt.buildSource(new CSVFileSourceV2<>(
                            FileUtils.pathJoin(this.snapshotPath, Constants.FILE_PERSON_OWN_ACCOUNT),
                            Collections.singletonList(Constants.PERSON_ID),
                            (CSVFileSourceV2.CSVLineParser<IVertex<Long, Double>>) line -> {
                                IVertex<Long, Double> vertex = new ValueVertex<>(
                                        Long.valueOf(line[0]),
                                        0.0
                                );
                                return Collections.singletonList(vertex);
                            }
                    ), AllWindow.getInstance())
                    .withParallelism(Constants.SOURCE_PARALLELISM);

            PWindowStream<IVertex<Long, Double>> unionVertexStream = accountVertices
                    .union(personVertices)
                    .withParallelism(Constants.UNION_PARALLELISM)
                    .union(loanVertices)
                    .withParallelism(Constants.UNION_PARALLELISM);

            PWindowSource<IEdge<Long, Double>> personOwnAccountEdges = pipelineTaskCxt.buildSource(new CSVFileSourceV2<>(
                            FileUtils.pathJoin(this.snapshotPath, Constants.FILE_PERSON_OWN_ACCOUNT),
                            Arrays.asList(Constants.PERSON_ID, Constants.ACCOUNT_ID),
                            (CSVFileSourceV2.CSVLineParser<IEdge<Long, Double>>) line -> {
                                IEdge<Long, Double> edge = new ValueEdge<>(
                                        Long.valueOf(line[1]),
                                        Long.valueOf(line[0]),
                                        null,
                                        EdgeDirection.IN
                                );
                                return Collections.singletonList(edge);
                            }
                    ), AllWindow.getInstance())
                    .withParallelism(Constants.SOURCE_PARALLELISM);

            PWindowSource<IEdge<Long, Double>> loanDepositAccountEdges = pipelineTaskCxt.buildSource(new CSVFileSourceV2<>(
                            FileUtils.pathJoin(this.snapshotPath, Constants.FILE_LOAN_DEPOSIT_ACCOUNT),
                            Arrays.asList(Constants.LOAN_ID, Constants.ACCOUNT_ID, Constants.AMOUNT),
                            (CSVFileSourceV2.CSVLineParser<IEdge<Long, Double>>) line -> {
                                IEdge<Long, Double> edge = new ValueEdge<>(
                                        -Long.parseLong(line[0]),
                                        Long.valueOf(line[1]),
                                        Double.valueOf(line[2])
                                );
                                return Collections.singletonList(edge);
                            }
                    ), AllWindow.getInstance())
                    .withParallelism(Constants.SOURCE_PARALLELISM);

            PWindowSource<IEdge<Long, Double>> accountTransferAccountEdges = pipelineTaskCxt.buildSource(new CSVFileSourceV2<>(
                            FileUtils.pathJoin(this.snapshotPath, Constants.FILE_ACCOUNT_TRANSFER_ACCOUNT),
                            Arrays.asList(Constants.FROM_ID, Constants.TO_ID, Constants.AMOUNT),
                            (CSVFileSourceV2.CSVLineParser<IEdge<Long, Double>>) line -> {
                                IEdge<Long, Double> edge = new ValueEdge<>(
                                        Long.valueOf(line[0]),
                                        Long.valueOf(line[1]),
                                        Double.valueOf(line[2])
                                );
                                return Collections.singletonList(edge);
                            }
                    ), AllWindow.getInstance())
                    .withParallelism(Constants.SOURCE_PARALLELISM);

            PWindowStream<IEdge<Long, Double>> unionEdgeStream = personOwnAccountEdges
                    .union(loanDepositAccountEdges)
                    .withParallelism(Constants.UNION_PARALLELISM)
                    .union(accountTransferAccountEdges)
                    .withParallelism(Constants.UNION_PARALLELISM);

            int iterationParallelism = Constants.ITERATOR_PARALLELISM;
            GraphViewDesc graphViewDesc = GraphViewBuilder
                    .createGraphView("case1")
                    .withShardNum(iterationParallelism)
                    .withBackend(IViewDesc.BackendType.Memory)
                    .build();

            PGraphWindow<Long, Double, Double> graphWindow =
                    pipelineTaskCxt.buildWindowStreamGraph(unionVertexStream, unionEdgeStream, graphViewDesc);

            graphWindow.compute(new Case1VCAlgorithm(4))
                    .compute(iterationParallelism)
                    .getVertices()
                    .filter((FilterFunction<IVertex<Long, Double>>) record -> record.getId() > 0 && record.getValue() != null && record.getValue() > 0)
                    .withParallelism(Constants.FILTER_PARALLELISM)
                    .map((MapFunction<IVertex<Long, Double>, LongStringResult>) value ->
                            new LongStringResult(value.getId(), String.format("%.2f", value.getValue() / 100_000_000)))
                    .withParallelism(Constants.MAP_PARALLELISM)
                    .sink(new CSVFileSinkV2<>("result1.csv"))
                    .withParallelism(Constants.SINK_PARALLELISM);
        });
    }

    public static class Case1VCAlgorithm extends VertexCentricCompute<Long, Double, Double, Tuple<Long, Double>> {

        public Case1VCAlgorithm(long iterations) {
            super(iterations);
        }

        @Override
        public VertexCentricCombineFunction<Tuple<Long, Double>> getCombineFunction() {
            return null;
        }

        @Override
        public VertexCentricComputeFunction<Long, Double, Double, Tuple<Long, Double>> getComputeFunction() {
            return new Case4VertexCentricComputeFunction();
        }
    }

    public static class Case4VertexCentricComputeFunction extends AbstractVcFunc<Long, Double, Double, Tuple<Long, Double>> {

        @Override
        public void compute(Long vertexId, Iterator<Tuple<Long, Double>> messageIterator) {
            switch ((int) context.getIterationId()) {
                case 1:
                    Double vertexValue = this.context.vertex().get().getValue();
                    if (vertexValue != null && vertexValue > 0) {
                        List<IEdge<Long, Double>> outEdges = this.context.edges().getOutEdges();
                        if (!outEdges.isEmpty()) {
                            Tuple<Long, Double> msg = Tuple.of(vertexId, vertexValue);
                            for (IEdge<Long, Double> outEdge : outEdges) {
                                this.context.sendMessage(outEdge.getTargetId(), msg);
                            }
                        }
                    }
                    break;
                case 2:
                    List<IEdge<Long, Double>> outEdges = this.context.edges().getOutEdges();
                    while (messageIterator.hasNext()) {
                        Tuple<Long, Double> msg = messageIterator.next();
                        for (IEdge<Long, Double> outEdge : outEdges) {
                            this.context.sendMessage(outEdge.getTargetId(), msg);
                        }
                    }
                    break;
                case 3:
                    List<IEdge<Long, Double>> inEdges = this.context.edges().getInEdges();
                    while (messageIterator.hasNext()) {
                        Tuple<Long, Double> msg = messageIterator.next();
                        for (IEdge<Long, Double> inEdge : inEdges) {
                            this.context.sendMessage(inEdge.getTargetId(), msg);
                        }
                    }
                    break;
                case 4:
                    double personLoanSum = 0.0;
                    Set<Long> loanSet = new HashSet<>();
                    while (messageIterator.hasNext()) {
                        Tuple<Long, Double> msg = messageIterator.next();
                        Long loanId = msg.f0;
                        if (!loanSet.contains(loanId)) {
                            personLoanSum += msg.f1;
                            loanSet.add(loanId);
                        }
                    }
                    this.context.setNewVertexValue(personLoanSum);
            }
        }
    }
}
